#pragma once
#include "NvInferVersion.h"
#include "NvInfer.h"
#include <vector>
#include <string>
#include <cassert>
#include <iostream>
#include <cuda_fp16.h> // 需要包含此头文件来支持 __half

namespace nvinfer1
{
  class HardSwishPlugin : public IPluginV2DynamicExt
  {
  public:
    HardSwishPlugin();
    HardSwishPlugin(const void *data, size_t length);
    ~HardSwishPlugin() override = default;

    int getNbOutputs() const noexcept override;
    DimsExprs getOutputDimensions(int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder) noexcept override;

    int initialize() noexcept override;
    virtual void terminate() noexcept override;

    // virtual size_t getWorkspaceSize(int maxBatchSize) const noexcept override;
    virtual size_t getWorkspaceSize(const PluginTensorDesc* inputs, int nbInputs, const PluginTensorDesc* outputs, int nbOutputs) const TRTNOEXCEPT override;

    // virtual int enqueue(int batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) noexcept override;
    virtual int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
    virtual size_t getSerializationSize() const noexcept override;

    virtual void serialize(void *buffer) const noexcept override;
    void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) noexcept override;
    // void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format) override;
    // 这个函数可以获取到数据类型和输入的维度信息，如果有需要用到的可以在这里将相关信息取出来
    // void configurePlugin(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs,
    //                      const DataType *inputTypes, const DataType *outputTypes, const bool *inputIsBroadcast,
    //                      const bool *outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) override;

    virtual void configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out, int nbOutputs) TRTNOEXCEPT override;


    // bool supportsFormatCombination(int pos, const PluginTensorDesc *inOut, int nbInputs, int nbOutputs) const override;
    // bool supportsFormat(DataType type, PluginFormat format) const override;

    bool supportsFormatCombination(int32_t pos, const nvinfer1::PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;

    const char *getPluginType() const noexcept override;
    const char *getPluginVersion() const noexcept override;
    void destroy() noexcept override;
    IPluginV2DynamicExt *clone() const noexcept override;

    // 添加缺失的函数声明
    DataType getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const noexcept override;
    // bool isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted, int nbInputs) const noexcept override;
    // bool canBroadcastInputAcrossBatch(int inputIndex) const noexcept override;
    void setPluginNamespace(const char *pluginNamespace) noexcept override;

    const char *getPluginNamespace() const noexcept override;
    
  private:
    std::string mPluginNamespace;
    Dims inputDims_;  // 存储输入维度
    
  };

  // 插件创建类：用于创建和管理插件，提供插件元信息（名称、版本等），创建插件实例、管理插件命名空间
  class HardSwishPluginCreator : public IPluginCreator
  {
  public:
    HardSwishPluginCreator(); // 构造函数需初始化插件属性（如 mPluginAttributes 和 mFC）
    ~HardSwishPluginCreator() override = default;
    const char *getPluginName() const override;            // 插件命令
    const char *getPluginVersion() const override;         // 插件版本
    const PluginFieldCollection *getFieldNames() override; // 返回插件的参数字段集合（mFC），描述插件支持的配置参数（如 alpha、beta等）
    // 根据字段集合 fc 创建插件实例（如解析参数并构造 HardSwishPlugin）,name是插件示例的名字，方便调试
    IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) override;
    // 从序列化数据中重建插件实例（如加载训练好的模型时调用）
    IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override;

    // 设置/获取插件的命名空间，用于避免插件名称冲突
    void setPluginNamespace(const char *pluginNamespace) override;
    const char *getPluginNamespace() const override;

  private:
    static PluginFieldCollection mFC;                  // 模型转换时，将其他模型参数列表传递给 TensorRT 的核心接口
    static std::vector<PluginField> mPluginAttributes; // 储存插件参数的定义，比如下方定义了插件有两个参数，以及类型和默认值

    std::string mPluginNamespace;
  };
};
